///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/fqmc/orb_opt_acc.cpp
///
/// \brief   implementation file for accumulator to collect pieces needed for orbital optimization
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<formic/fqmc/orb_opt_acc.h>

#include<numeric>
#include<cassert>
#include<algorithm>
#include<cmath>
#include<complex>

#include<boost/format.hpp>
#include<boost/shared_ptr.hpp>

#include<formic/exception.h>
#include<formic/input/base.h>
#include<formic/mpi/interface.h>
#include<formic/numeric/numeric.h>
#include<formic/output/output.h>
#include<formic/lapack/interface.h>
#include<formic/wfn/configuration.h>
#include<formic/wfn/coeff_obj.h>
#include<formic/fqmc/accumulator.h>
#include<formic/vector/vector.h>
#include<formic/ham/tei.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  constructor
///
/// \param[in]      norb     number of orbitals in basis
/// \param[in]      oc_mat   the orbital coefficient matrix for the current orbital rotation
/// \param[in]      co       coefficient object for the wave function
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
formic::fqmc::OrbOptAcc<S>::OrbOptAcc(const int norb,
                                      const formic::ConstMatrix<S,int> & oc_mat,
                                      const boost::shared_ptr<formic::CoeffObj<S> > & gf,
                                      const boost::shared_ptr<formic::CoeffObj<S> > & vf)
  : formic::Accumulator<S>(gf, vf),
    _norb(norb),
    _oei_der(),
    _tei_dns_der(),
    _tei_cmp_der(),
    _oc(oc_mat),
    //_bare_oc_der_rat(),
    //_enrg_oc_der_rat(),
    //_locl_oc_der_rat(),
    _energy(formic::zero(S()))
{}

template formic::fqmc::OrbOptAcc<double>::OrbOptAcc(const int,
                                                    const formic::ConstMatrix<double,int> &,
                                                    const boost::shared_ptr<formic::CoeffObj<double> > &,
                                                    const boost::shared_ptr<formic::CoeffObj<double> > &);
template formic::fqmc::OrbOptAcc<std::complex<double> >::OrbOptAcc(const int,
                                                                   const formic::ConstMatrix<std::complex<double>,int> &,
                                                                   const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > &,
                                                                   const boost::shared_ptr<formic::CoeffObj<std::complex<double> > > &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Accumulate the local contributions integral derivatives
///
/// \param[in]      userinp  the user's input options
/// \param[in]      w        the weight to apply to this configuration's contribution
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::fqmc::OrbOptAcc<S>::child_accumulate(const formic::InputBase & userinp, const double w) {

//  // get the square of the value/guiding ratio
//  const double vg_square = ( _vf == _gf  ?  1.0  :  formic::square_norm( ( _vf->get_coeff() / _gf->get_coeff() ).get() ) );
//
//  // accumulate the squares of the value/guiding ratios
//  _vgs_total += w * vg_square;

  // get a variable to hold the local energy
  S loc_en = formic::zero(S());

  // accumulate the derivative adjoints of the modified one electron integrals and the two electron integrals
  this->_base_vf->local_energy_integral_der(w * this->_vgs, loc_en, &_oei_der.at(0,0), &_tei_dns_der.at(0));

  // update the total local energy
  _energy += w * this->_vgs * loc_en;

}

template void formic::fqmc::OrbOptAcc<double>::child_accumulate(const formic::InputBase & userinp, const double w);
template void formic::fqmc::OrbOptAcc<std::complex<double> >::child_accumulate(const formic::InputBase & userinp, const double w);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief finalizes the accumulation (this may include for example communicating the results)
///
/// \param[in]      userinp  the user's input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::fqmc::OrbOptAcc<S>::child_finalize(const formic::InputBase & userinp) {

//  // reduce sum of squares of the value/guiding ratios
//  {
//    double vgs_copy = _vgs_total / this->_total_weight;
//    formic::mpi::allreduce(&vgs_copy, &_vgs_total, 1, MPI::SUM);
//  }

  // reduce energy to all processes and scale it by the total weight
  {
    S e_copy = _energy / this->_total_weight / this->_vgsa;
    formic::mpi::allreduce(&e_copy, &_energy, 1, MPI::SUM);
  }

  //if ( formic::mpi::rank() == 0 )
  //  formic::of << boost::format("formic::fqmc::OrbOptAcc<S> final energy = %20.12f") % _energy << std::endl;

//  // compute component of energy derivatives w.r.t. orbital coefficients when the derivative hits the wave function
//  // (derivatives hitting the hamiltonian will be handled later using the integral derivatives)
//  _enrg_oc_der_rat.axpy(-_energy, _bare_oc_der_rat);
//  _enrg_oc_der_rat *= 2.0 * formic::unity(S()) / this->_total_weight;
//  formic::mpi::reduce(_enrg_oc_der_rat);

//  // bare and local derivative ratio vectors can be discarded now
//  _bare_oc_der_rat.clear();
//  _locl_oc_der_rat.clear();

  // reduce accumulated integrals to root process
  formic::mpi::reduce(_oei_der);
  formic::mpi::reduce(_tei_dns_der);

  // root process scales reduced integrals by the total weight
  if ( formic::mpi::rank() == 0 ) {
    _oei_der /= formic::unity(S()) * this->_total_weight * this->_vgsa;
    _tei_dns_der /= formic::unity(S()) * this->_total_weight * this->_vgsa;
  }

  // root process places two electron integrals in compact array and discards the dense array
  if ( formic::mpi::rank() == 0 ) {
    _tei_cmp_der.reset(formic::tei_compact_array_length(_norb));
    formic::tei_dense_to_compact(_norb, &_tei_dns_der.at(0), &_tei_cmp_der.at(0));
    _tei_dns_der.clear();
  }

  // root process adds modified one-electron integral contribution to the two electron integral derivatives
  if ( formic::mpi::rank() == 0 ) {
    for (int i = 0; i < _norb; i++)
    for (int j = 0; j < _norb; j++)
    for (int p = 0; p < _norb; p++) {
      double x = 1.0;
      x *= ( i == p ? 1.0 : 0.5 );
      x *= ( j == p ? 1.0 : 0.5 );
      x *= ( std::min(i,p) == std::min(j,p) && std::max(i,p) == std::max(j,p) ? 1.0 : 0.5 );
      _tei_cmp_der[ formic::tei_compound_index(i,p,p,j) ] -= 0.5 * _oei_der(i,j) * x;
      //_mei[i*_norb+j] -= 0.5 * this->get_tei(i,p,p,j);
    }
  }

  // non-root processes clear their storage
  if ( formic::mpi::rank() != 0 ) {
    _oei_der.clear();
    _tei_dns_der.clear();
    _tei_cmp_der.clear();
    //_enrg_oc_der_rat.clear();
    //_bare_oc_der_rat.clear();
    //_locl_oc_der_rat.clear();
  }

}

template void formic::fqmc::OrbOptAcc<double>::child_finalize(const formic::InputBase & userinp);
template void formic::fqmc::OrbOptAcc<std::complex<double> >::child_finalize(const formic::InputBase & userinp);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  resets the accumulator
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class S>
void formic::fqmc::OrbOptAcc<S>::child_reset(const formic::Configuration & config) {

  //_bare_oc_der_rat.reset(_norb, _norb, formic::zero(S()));
  //_enrg_oc_der_rat.reset(_norb, _norb, formic::zero(S()));
  //_locl_oc_der_rat.reset(_norb, _norb, formic::zero(S()));

  _oei_der.reset(_norb, _norb, formic::zero(S()));
  _tei_dns_der.reset(_norb*_norb*_norb*_norb, formic::zero(S()));

  _tei_cmp_der.clear();

  _energy = formic::zero(S());

  //_vgs_total = 0.0;

}

template void formic::fqmc::OrbOptAcc<double>::child_reset(const formic::Configuration & config);
template void formic::fqmc::OrbOptAcc<std::complex<double> >::child_reset(const formic::Configuration & config);
